package com.dms_back1.interceptor;

import com.dms_back1.annotation.RateLimit;
import com.dms_back1.ex.BusinessException;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Slf4j
@Aspect
@Component
public class RateLimitInterceptor {

    private final StringRedisTemplate redisTemplate;
    private final HttpServletRequest request;
    private final ConcurrentHashMap<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    public RateLimitInterceptor(StringRedisTemplate redisTemplate, HttpServletRequest request) {
        this.redisTemplate = redisTemplate;
        this.request = request;
    }

    @Around("@annotation(com.dms_back1.annotation.RateLimit)")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);

        String key = generateKey(rateLimit, request);
        String countKey = "rate_limit:" + key;

        // 检查Redis中的计数器
        String count = redisTemplate.opsForValue().get(countKey);
        if (count == null) {
            // 设置计数器和过期时间
            redisTemplate.opsForValue().set(countKey, "1", rateLimit.time(), rateLimit.unit());
        } else {
            long currentCount = Long.parseLong(count);
            if (currentCount >= rateLimit.count()) {
                log.warn("请求被限流: key={}, count={}", key, currentCount);
                throw new BusinessException(rateLimit.message());
            }
            // 增加计数器
            redisTemplate.opsForValue().increment(countKey);
        }

        // 使用Guava RateLimiter进行单机限流
        RateLimiter rateLimiter = rateLimiterMap.computeIfAbsent(key,
            k -> RateLimiter.create(rateLimit.count() / (double) rateLimit.unit().toSeconds(rateLimit.time())));

        if (!rateLimiter.tryAcquire(100, TimeUnit.MILLISECONDS)) {
            log.warn("请求被限流(单机): key={}", key);
            throw new BusinessException(rateLimit.message());
        }

        return point.proceed();
    }

    private String generateKey(RateLimit rateLimit, HttpServletRequest request) {
        String key = rateLimit.key();
        if (key.isEmpty()) {
            key = request.getRequestURI();
        }
        String ip = getClientIp(request);
        return key + ":" + ip;
    }

    private String getClientIp(HttpServletRequest request) {
        String ip = request.getHeader("X-Forwarded-For");
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_CLIENT_IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_X_FORWARDED_FOR");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
} 